import os
import cv2
import numpy as np

import scipy.io
import os

"""
适用于将ITCVD的mat格式数据转为yolo格式的txt数据
"""
# Define class mapping
class_mapping = {
    'object': 0,  # Assuming all objects have the same class
}

# Function to convert bounding box to YOLO format
def convert_to_yolo(bbox, img_width, img_height):
    x_center = (bbox[0] + bbox[2]) / (2 * img_width)
    y_center = (bbox[1] + bbox[3]) / (2 * img_height)
    width = (bbox[2] - bbox[0]) / img_width
    height = (bbox[3] - bbox[1]) / img_height
    return x_center, y_center, width, height

# Path to the directory containing images
images_directory = 'D:/datasets/ITCVD/Testing/Image'  # Update with the directory where your images are stored

# Loop through each image in the directory
for filename in os.listdir(images_directory):
    if filename.endswith(('.jpg', '.jpeg', '.png')):  # Check if the file is an image file
        # Read the image
        image_path = os.path.join(images_directory, filename)
        img = cv2.imread(image_path)
        img_height, img_width, _ = img.shape

        # Accessing the object bounding box information for the current image
        gt_mat_path = 'D:/datasets/ITCVD/Testing/GT'
        mat_name = filename[:-4] + ".mat"
        mat_path = os.path.join(gt_mat_path,mat_name)
        print(mat_path)

        itcvd_data = scipy.io.loadmat(mat_path)

        object_bboxes = itcvd_data['x' + filename[:-4]]

        # Create YOLO format annotation file for the current image
        with open(os.path.splitext(image_path)[0] + '.txt', 'w') as f:
            for bbox in object_bboxes:
                x_min, y_min, x_max, y_max = bbox[0], bbox[1], bbox[2], bbox[3]
                yolo_bbox = convert_to_yolo([x_min, y_min, x_max, y_max], img_width, img_height)
                class_id = class_mapping['object']
                yolo_line = f"{class_id} {' '.join(str(coord)[:8] for coord in yolo_bbox)}"
                f.write(yolo_line + '\n')
